import tensorflow as tf
from data_utils import load_train32C, load_val32C, gaussian_noise
import numpy as np
from models import FCN
import os
import cv2


class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def train(num_epochs=10, batch_size=32):
    train_imgs, train_names = load_train32C(Debug=False)
    val_imgs, val_names = load_val32C(Debug=False)
    train_X = np.asarray(train_imgs, np.uint8)
    val_X = np.asarray(val_imgs, np.uint8)
    train_X_noised = gaussian_noise(train_X)
    val_X_noised = gaussian_noise(val_X)
    train_X = np.asarray(train_X, np.uint8) / 255.
    val_X = np.asarray(val_X, np.uint8) / 255.

    sess = tf.Session()
    with tf.device('/gpu:1'):
        model = FCN(batch_size=batch_size, img_size=32)
        tf.initialize_all_variables().run(session=sess)
    saver = tf.train.Saver(tf.all_variables(), max_to_keep=100)
    tf_writer = tf.train.SummaryWriter(
        '/media/wjsun/delldisk/dell/wxm/Data/decsai/CM512/dip_exs/wri/train_g', sess.graph)
    tf_writer_val = tf.train.SummaryWriter(
        '/media/wjsun/delldisk/dell/wxm/Data/decsai/CM512/dip_exs/wri/val_g', sess.graph)

    nb_iters = 0
    nb_iters_val = 0
    for e in range(num_epochs):
        for start, end in zip(range(0, len(train_X), batch_size),
                              range(batch_size, len(train_X), batch_size)):
            nb_iters += 1
            result = sess.run([model.loss, model.merged, model.train_op], {model.input_imgs: train_X_noised[start:end],
                                                                           model.label: train_X[start:end]})
            # result[0] /= (420. * 580. * batch_size)
            tf_writer.add_summary(result[1], nb_iters)
            if nb_iters % 200 == 0:
                print '{0} train samples'.format(len(train_X))
                print '{0} val samples'.format(len(val_X))
                print 'total epochs: {0}'.format(num_epochs)
                print '-' * 30
                print 'Epoch: {0} | nb_iters: {1}'.format(e, nb_iters)
                print bcolors.OKBLUE + 'loss: {0}'.format(result[0]) + bcolors.ENDC
                print '*' * 50
        for start, end in zip(range(0, len(val_X), batch_size),
                              range(batch_size, len(val_X), batch_size)):
            result_val = sess.run([model.loss, model.merged],
                                  feed_dict={model.input_imgs: val_X_noised[start: end],
                                             model.label: val_X[start:end]})
            # result_val[0] /= (420. * 580. * batch_size)
            if nb_iters_val % 200 == 0:
                print 'Epoch: {0}'.format(e)
                print bcolors.OKBLUE + 'val_cost: {0}'.format(result_val[0]) + bcolors.ENDC
                print '*' * 50
            tf_writer_val.add_summary(result_val[1], nb_iters_val)
            nb_iters_val += 1
        cp_path = '/media/wjsun/delldisk/dell/wxm/Data/decsai/CM512/dip_exs/saver/cpgaussian/'
        if not os.path.exists(cp_path):
            os.mkdir(cp_path)
        saver.save(sess, cp_path, global_step=e)
        print 'model in epoch {0} saved successfully'.format(e)
        print '*' * 50


def inf():
    sess = tf.Session()
    with tf.device('/gpu:1'):
        model = FCN(batch_size=32, img_size=512)
        tf.initialize_all_variables().run(session=sess)
    saver = tf.train.Saver(tf.all_variables())
    model_path = '/media/wjsun/delldisk/dell/wxm/Data/decsai/CM512/dip_exs/saver/cpgaussian/-0'
    saver.restore(sess, model_path)

    img1 = cv2.imread('gaussian/lena_gaussian.png', 0)
    img2 = cv2.imread('gaussian/barbara_gaussian.png', 0)
    imgs = []
    imgs.append(img1)
    imgs.append(img2)
    imgs = np.asarray(imgs, np.float32)
    imgs /= 255.
    paddings = np.random.randint(-100, 100, [30, 512, 512]) / 100.
    img_arrs_concat = np.concatenate((imgs, paddings), axis=0)

    result_fusion = sess.run(model.den, {model.input_imgs: img_arrs_concat,
                                         model.label: img_arrs_concat})
    result_fusion = np.asarray(result_fusion, np.uint8)
    cv2.imwrite('cnnRes/lena.png', result_fusion[0])
    cv2.imwrite('cnnRes/barbara.png', result_fusion[1])


inf()
